from mpi4py import MPI
from enthought.tvtk.api import tvtk,write_data
from numpy import c_,ravel,empty,savez,array

def write_vtkxml_sgrid(data, filename):
    if type(filename) == type(''):
        file = open(filename,'w')
    for k in data.keys():
        data[k] = data[k].astype('f')
    x,y,z = [data[i] for i in ('x','y','z')]
    dims = x.shape
    var = [i for i in data.keys() if i not in ('x','y','z')]
    sg = tvtk.StructuredGrid(points=c_[x.flat,y.flat,z.flat],dimensions=array(dims))
    pd = tvtk.PointData()
    for v in var:
        pd.scalars = ravel(data[v].astype('f'))
        pd.scalars.name = v
        sg.point_data.add_array(pd.scalars)
    if 'u' in var and 'v' in var and 'w' in var:
        print 'Detected velocity vector'
        u,v,w = [data[i] for i in ('u','v','w')]
        sg.point_data.vectors = c_[ravel(u),ravel(v),ravel(w)]
        sg.point_data.vectors.name = 'V'
    xsgw = tvtk.XMLStructuredGridWriter(input=sg, file_name=filename)
    xsgw.write()

def write(elems,filename,isgrid=False,gridsize=None):
    """
    Writes position, vorticity and velocity data to appropriate files
    """
    comm = MPI.COMM_WORLD
    size = len(elems)
    x,y,z,gx,gy,gz,vx,vy,vz = (empty(size),empty(size),empty(size),
                               empty(size),empty(size),empty(size),
                               empty(size),empty(size),empty(size))
    for i,elem in enumerate(elems):
        x[i] = elem.pos.x
        y[i] = elem.pos.y
        z[i] = elem.pos.z
        gx[i] = elem.vor.x
        gy[i] = elem.vor.y
        gz[i] = elem.vor.z
        vx[i] = elem.vel.x
        vy[i] = elem.vel.y
        vz[i] = elem.vel.z
    
    if isgrid: # need to save velocity field data in structured grid format
        num = (gridsize,gridsize,gridsize)
        xg3d,yg3d,zg3d,vxg3d,vyg3d,vzg3d = (empty(num),empty(num),empty(num),
                                            empty(num),empty(num),empty(num))
        xg3d.flat[:] = x
        yg3d.flat[:] = y
        zg3d.flat[:] = z
        vxg3d.flat[:] = vx
        vyg3d.flat[:] = vy
        vzg3d.flat[:] = vz
        data_g = {'x':xg3d,'y':yg3d,'z':zg3d,'u':vxg3d,'v':vyg3d,'w':vzg3d}
        write_vtkxml_sgrid(data_g,filename+'.vts')
        savez(filename+'.npz', **data_g)
    
    else:
        # add point data
        pd = tvtk.PolyData(points=c_[x,y,z])

        # add vor data to points
        vors = c_[gx,gy,gz]
        pvors = tvtk.DoubleArray()
        pvors.name = 'vor'
        pvors.from_array(vors)
        pd.point_data.add_array(pvors)
        
        # add vel data to points
        vels = c_[vx,vy,vz]
        pvels = tvtk.DoubleArray()
        pvels.name = 'vel'
        pvels.from_array(vels)
        pd.point_data.add_array(pvels)
        
        # save point data
        write_data(pd,filename+'.vtk') # tvtk's write_data
        data_b = {'x':x,'y':y,'z':z,'gx':gx,'gy':gy,'gz':gz,'vx':vx,'vy':vy,'vz':vz}
        savez(filename+'.npz',**data_b)
